"""
坐姿分析算法模块
"""

import numpy as np
from typing import List, Tuple, Dict, Optional
from dataclasses import dataclass
import json
import os
from datetime import datetime
from config import (POSTURE_THRESHOLDS, SCORING_WEIGHTS, POSE_LANDMARKS, 
                   POSTURE_SUGGESTIONS)
from utils import (calculate_head_desk_distance, calculate_neck_posture_comprehensive,
                  calculate_shoulder_posture_detailed, calculate_spine_alignment_advanced,
                  calculate_sitting_stability, validate_landmarks, get_posture_grade)

@dataclass
class PostureMetrics:
    """坐姿指标数据类（精细化版本）"""
    # 头部与桌面距离
    head_desk_distance: float = 0.0
    
    # 脊柱对齐
    spine_alignment_score: float = 0.0
    cervical_angle: float = 0.0
    thoracic_angle: float = 0.0
    
    # 颈椎姿态
    neck_forward_angle: float = 0.0
    neck_side_tilt: float = 0.0
    neck_twist_ratio: float = 0.0
    
    # 肩膀姿态
    shoulder_height_diff: float = 0.0
    shoulder_front_back_diff: float = 0.0
    shoulder_narrowing: float = 0.0
    shoulder_symmetry: float = 0.0
    
    # 坐姿稳定性
    sitting_stability: float = 0.0
    
    # 综合评估
    overall_score: float = 0.0
    grade: str = "poor"
    suggestions: List[str] = None
    
    def __post_init__(self):
        if self.suggestions is None:
            self.suggestions = []

class UserProfile:
    """用户档案类，用于个性化调整"""
    
    def __init__(self, user_id: str = "default"):
        self.user_id = user_id
        self.profile_file = f"user_profiles/{user_id}_profile.json"
        self.load_profile()
    
    def load_profile(self):
        """加载用户档案"""
        if os.path.exists(self.profile_file):
            try:
                with open(self.profile_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    self.age = data.get('age', 12)
                    self.height = data.get('height', 150)
                    self.learning_style = data.get('learning_style', 'normal')
                    self.posture_history = data.get('posture_history', [])
                    self.custom_thresholds = data.get('custom_thresholds', {})
                    self.preference_weights = data.get('preference_weights', {})
            except:
                self.init_default_profile()
        else:
            self.init_default_profile()
    
    def init_default_profile(self):
        """初始化默认档案"""
        self.age = 12
        self.height = 150
        self.learning_style = 'normal'
        self.posture_history = []
        self.custom_thresholds = {}
        self.preference_weights = {}
        
        # 创建目录
        os.makedirs(os.path.dirname(self.profile_file), exist_ok=True)
        self.save_profile()
    
    def save_profile(self):
        """保存用户档案"""
        try:
            data = {
                'age': self.age,
                'height': self.height,
                'learning_style': self.learning_style,
                'posture_history': self.posture_history[-100:],  # 只保存最近100条记录
                'custom_thresholds': self.custom_thresholds,
                'preference_weights': self.preference_weights,
                'last_updated': datetime.now().isoformat()
            }
            
            with open(self.profile_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"保存用户档案失败: {e}")
    
    def add_posture_record(self, metrics: PostureMetrics):
        """添加坐姿记录"""
        record = {
            'timestamp': datetime.now().isoformat(),
            'overall_score': metrics.overall_score,
            'grade': metrics.grade,
            'metrics': {
                'head_desk_distance': metrics.head_desk_distance,
                'spine_alignment_score': metrics.spine_alignment_score,
                'neck_forward_angle': metrics.neck_forward_angle,
                'neck_side_tilt': metrics.neck_side_tilt,
                'shoulder_height_diff': metrics.shoulder_height_diff,
                'shoulder_symmetry': metrics.shoulder_symmetry,
                'sitting_stability': metrics.sitting_stability
            }
        }
        
        self.posture_history.append(record)
        
        # 定期保存
        if len(self.posture_history) % 10 == 0:
            self.save_profile()
    
    def get_adaptive_thresholds(self) -> Dict:
        """获取自适应阈值"""
        if not self.posture_history:
            return {}
        
        # 分析用户历史数据
        recent_scores = [record['overall_score'] for record in self.posture_history[-50:]]
        avg_score = np.mean(recent_scores)
        
        # 根据用户平均表现调整阈值
        adaptive_thresholds = {}
        
        if avg_score > 85:  # 用户坐姿较好，可以提高标准
            adaptive_thresholds['strictness_factor'] = 1.2
        elif avg_score < 60:  # 用户坐姿较差，适当放宽标准
            adaptive_thresholds['strictness_factor'] = 0.8
        else:
            adaptive_thresholds['strictness_factor'] = 1.0
        
        return adaptive_thresholds
    
    def get_personalized_weights(self) -> Dict:
        """获取个性化权重"""
        # 基于年龄调整权重
        age_weights = {}
        
        if self.age < 10:  # 儿童，重点关注基础姿态
            age_weights = {
                'head_desk_distance': 0.25,
                'spine_alignment': 0.30,
                'neck_forward_angle': 0.20,
                'neck_side_tilt': 0.10,
                'neck_twist': 0.05,
                'shoulder_height_diff': 0.05,
                'shoulder_front_back_diff': 0.03,
                'shoulder_symmetry': 0.02,
                'sitting_stability': 0.00
            }
        elif self.age > 16:  # 青少年，重点关注精细姿态
            age_weights = {
                'head_desk_distance': 0.15,
                'spine_alignment': 0.20,
                'neck_forward_angle': 0.15,
                'neck_side_tilt': 0.12,
                'neck_twist': 0.08,
                'shoulder_height_diff': 0.12,
                'shoulder_front_back_diff': 0.08,
                'shoulder_symmetry': 0.08,
                'sitting_stability': 0.02
            }
        else:  # 默认权重
            age_weights = SCORING_WEIGHTS.copy()
        
        # 合并用户偏好权重
        final_weights = age_weights.copy()
        final_weights.update(self.preference_weights)
        
        return final_weights

class PostureAnalyzer:
    """坐姿分析器（精细化版本）"""
    
    def __init__(self, user_id: str = "default"):
        """初始化坐姿分析器"""
        self.thresholds = POSTURE_THRESHOLDS
        self.weights = SCORING_WEIGHTS
        self.landmarks_map = POSE_LANDMARKS
        self.suggestions = POSTURE_SUGGESTIONS
        
        # 用户档案
        self.user_profile = UserProfile(user_id)
        
        # 评分历史（用于平滑）
        self.score_history = []
        self.max_history_size = 10
        
        # 自适应学习参数
        self.learning_enabled = True
        self.adaptation_rate = 0.1
        
        print(f"坐姿分析器已初始化，用户ID: {user_id}")
        print(f"已加载用户档案，年龄: {self.user_profile.age}, 身高: {self.user_profile.height}cm")
        
    def analyze_posture(self, landmarks: List[Tuple[float, float]], 
                       image_width: int, image_height: int) -> PostureMetrics:
        """
        分析坐姿并返回详细指标
        
        Args:
            landmarks: 关键点坐标列表
            image_width: 图像宽度
            image_height: 图像高度
        
        Returns:
            PostureMetrics: 坐姿指标对象
        """
        # 验证关键点
        required_points = [0, 7, 8, 11, 12, 13, 14, 23, 24]  # 鼻子、耳朵、肩膀、肘部、髋部
        if not validate_landmarks(landmarks, required_points):
            return PostureMetrics()
        
        # 计算各项指标
        metrics = PostureMetrics()
        
        # 1. 头部与桌面距离
        metrics.head_desk_distance = calculate_head_desk_distance(landmarks, image_height)
        
        # 2. 脊柱对齐评估
        metrics.cervical_angle, metrics.thoracic_angle, metrics.spine_alignment_score = \
            calculate_spine_alignment_advanced(landmarks)
        
        # 3. 颈椎姿态评估
        metrics.neck_forward_angle, metrics.neck_side_tilt, metrics.neck_twist_ratio = \
            calculate_neck_posture_comprehensive(landmarks)
        
        # 4. 肩膀姿态评估
        metrics.shoulder_height_diff, metrics.shoulder_front_back_diff, \
        metrics.shoulder_narrowing, metrics.shoulder_symmetry = \
            calculate_shoulder_posture_detailed(landmarks)
        
        # 5. 坐姿稳定性评估
        metrics.sitting_stability = calculate_sitting_stability(landmarks)
        
        # 6. 计算综合评分（使用个性化权重）
        metrics.overall_score = self._calculate_overall_score_personalized(metrics)
        
        # 7. 平滑评分
        self.score_history.append(metrics.overall_score)
        if len(self.score_history) > self.max_history_size:
            self.score_history.pop(0)
        
        smoothed_score = np.mean(self.score_history)
        metrics.overall_score = smoothed_score
        
        # 8. 确定等级（使用自适应阈值）
        metrics.grade = self._get_adaptive_grade(metrics.overall_score)
        
        # 9. 生成个性化建议
        metrics.suggestions = self._generate_personalized_suggestions(metrics)
        
        # 10. 记录到用户档案
        if self.learning_enabled:
            self.user_profile.add_posture_record(metrics)
        
        return metrics
    
    def _calculate_overall_score_personalized(self, metrics: PostureMetrics) -> float:
        """
        使用个性化权重计算综合评分
        
        Args:
            metrics: 坐姿指标对象
        
        Returns:
            综合评分 (0-100)
        """
        # 获取个性化权重
        personalized_weights = self.user_profile.get_personalized_weights()
        
        # 各项指标的评分
        scores = {}
        
        # 头部与桌面距离评分
        scores['head_desk_distance'] = metrics.head_desk_distance
        
        # 脊柱对齐评分
        scores['spine_alignment'] = metrics.spine_alignment_score
        
        # 颈椎姿态评分
        scores['neck_forward_angle'] = self._evaluate_neck_forward(metrics.neck_forward_angle)
        scores['neck_side_tilt'] = self._evaluate_neck_side_tilt(metrics.neck_side_tilt)
        scores['neck_twist'] = self._evaluate_neck_twist(metrics.neck_twist_ratio)
        
        # 肩膀姿态评分
        scores['shoulder_height_diff'] = self._evaluate_shoulder_height_diff(metrics.shoulder_height_diff)
        scores['shoulder_front_back_diff'] = self._evaluate_shoulder_front_back_diff(metrics.shoulder_front_back_diff)
        scores['shoulder_symmetry'] = metrics.shoulder_symmetry
        
        # 坐姿稳定性评分
        scores['sitting_stability'] = metrics.sitting_stability
        
        # 加权计算综合评分
        total_score = 0.0
        for key, weight in personalized_weights.items():
            if key in scores:
                total_score += scores[key] * weight
        
        return min(100.0, max(0.0, total_score))
    
    def _evaluate_neck_forward(self, angle: float) -> float:
        """评估脖子前倾角度"""
        thresholds = self.thresholds['neck_posture']['forward_angle']
        
        if thresholds['excellent'][0] <= angle <= thresholds['excellent'][1]:
            return 100
        elif thresholds['good'][0] <= angle <= thresholds['good'][1]:
            return 80
        elif thresholds['fair'][0] <= angle <= thresholds['fair'][1]:
            return 60
        else:
            return 40
    
    def _evaluate_neck_side_tilt(self, angle: float) -> float:
        """评估脖子侧倾角度"""
        thresholds = self.thresholds['neck_posture']['side_tilt']
        
        if thresholds['excellent'][0] <= angle <= thresholds['excellent'][1]:
            return 100
        elif thresholds['good'][0] <= angle <= thresholds['good'][1]:
            return 80
        elif thresholds['fair'][0] <= angle <= thresholds['fair'][1]:
            return 60
        else:
            return 40
    
    def _evaluate_neck_twist(self, ratio: float) -> float:
        """评估脖子扭转程度"""
        thresholds = self.thresholds['neck_posture']['twist_ratio']
        
        if thresholds['excellent'][0] <= ratio <= thresholds['excellent'][1]:
            return 100
        elif thresholds['good'][0] <= ratio <= thresholds['good'][1]:
            return 80
        elif thresholds['fair'][0] <= ratio <= thresholds['fair'][1]:
            return 60
        else:
            return 40
    
    def _evaluate_shoulder_height_diff(self, angle: float) -> float:
        """评估肩膀高低差"""
        thresholds = self.thresholds['shoulder_posture']['height_diff']
        
        if thresholds['excellent'][0] <= angle <= thresholds['excellent'][1]:
            return 100
        elif thresholds['good'][0] <= angle <= thresholds['good'][1]:
            return 80
        elif thresholds['fair'][0] <= angle <= thresholds['fair'][1]:
            return 60
        else:
            return 40
    
    def _evaluate_shoulder_front_back_diff(self, diff: float) -> float:
        """评估肩膀前后差"""
        thresholds = self.thresholds['shoulder_posture']['front_back_diff']
        
        if thresholds['excellent'][0] <= diff <= thresholds['excellent'][1]:
            return 100
        elif thresholds['good'][0] <= diff <= thresholds['good'][1]:
            return 80
        elif thresholds['fair'][0] <= diff <= thresholds['fair'][1]:
            return 60
        else:
            return 40
    
    def _generate_suggestions(self, metrics: PostureMetrics) -> List[str]:
        """
        生成坐姿改进建议
        
        Args:
            metrics: 坐姿指标对象
        
        Returns:
            建议列表
        """
        suggestions = []
        
        # 头部与桌面距离建议
        if metrics.head_desk_distance < 60:
            if metrics.head_desk_distance < 40:
                suggestions.append(self.suggestions['head_distance']['too_close'])
            else:
                suggestions.append(self.suggestions['head_distance']['too_far'])
        
        # 脊柱对齐建议
        if metrics.spine_alignment_score < 70:
            if metrics.spine_alignment_score < 50:
                suggestions.append(self.suggestions['spine_alignment']['poor'])
            else:
                suggestions.append(self.suggestions['spine_alignment']['fair'])
        
        # 颈椎姿态建议
        if metrics.neck_forward_angle > 20:
            suggestions.append(self.suggestions['neck_posture']['forward_tilt'])
        
        if metrics.neck_side_tilt > 15:
            suggestions.append(self.suggestions['neck_posture']['side_tilt'])
        
        if metrics.neck_twist_ratio > 0.2:
            suggestions.append(self.suggestions['neck_posture']['twist'])
        
        # 肩膀姿态建议
        if metrics.shoulder_height_diff > 8:
            suggestions.append(self.suggestions['shoulder_posture']['height_diff'])
        
        if metrics.shoulder_front_back_diff > 25:
            suggestions.append(self.suggestions['shoulder_posture']['front_back_diff'])
        
        if metrics.shoulder_narrowing > 0.05:
            suggestions.append(self.suggestions['shoulder_posture']['narrowing'])
        
        if metrics.shoulder_symmetry < 60:
            suggestions.append(self.suggestions['shoulder_posture']['asymmetry'])
        
        # 坐姿稳定性建议
        if metrics.sitting_stability < 70:
            suggestions.append(self.suggestions['sitting_stability']['unstable'])
        
        # 如果没有具体建议，给出总体建议
        if not suggestions:
            suggestions.append(self.suggestions['general'][metrics.grade])
        
        return suggestions
    
    def _generate_personalized_suggestions(self, metrics: PostureMetrics) -> List[str]:
        """
        生成个性化建议
        
        Args:
            metrics: 坐姿指标对象
        
        Returns:
            建议列表
        """
        suggestions = []
        
        # 基础建议生成
        base_suggestions = self._generate_suggestions(metrics)
        suggestions.extend(base_suggestions)
        
        # 基于年龄的个性化建议
        age = self.user_profile.age
        if age < 10:
            suggestions.append("小朋友，记得保持背部挺直，就像小树苗一样！")
        elif age < 16:
            suggestions.append("青少年正处于发育期，良好的坐姿对成长很重要哦！")
        else:
            suggestions.append("长期保持良好坐姿，有助于预防脊柱问题。")
        
        # 基于历史数据的建议
        if len(self.user_profile.posture_history) > 10:
            recent_scores = [record['overall_score'] for record in self.user_profile.posture_history[-10:]]
            trend = np.mean(recent_scores[-5:]) - np.mean(recent_scores[:5])
            
            if trend > 5:
                suggestions.append("最近的坐姿有明显改善，继续保持！")
            elif trend < -5:
                suggestions.append("最近的坐姿有所下降，请注意调整。")
        
        return suggestions[:5]  # 最多返回5条建议
    
    def get_detailed_analysis(self, metrics: PostureMetrics) -> Dict[str, any]:
        """
        获取详细的分析报告
        
        Args:
            metrics: 坐姿指标对象
        
        Returns:
            详细分析报告
        """
        return {
            'overall_score': metrics.overall_score,
            'grade': metrics.grade,
            'detailed_metrics': {
                'head_desk_distance': {
                    'value': metrics.head_desk_distance,
                    'status': self._get_metric_status(metrics.head_desk_distance, 'head_distance')
                },
                'spine_alignment': {
                    'score': metrics.spine_alignment_score,
                    'cervical_angle': metrics.cervical_angle,
                    'thoracic_angle': metrics.thoracic_angle,
                    'status': self._get_metric_status(metrics.spine_alignment_score, 'spine_alignment')
                },
                'neck_posture': {
                    'forward_angle': metrics.neck_forward_angle,
                    'side_tilt': metrics.neck_side_tilt,
                    'twist_ratio': metrics.neck_twist_ratio,
                    'status': self._get_neck_status(metrics)
                },
                'shoulder_posture': {
                    'height_diff': metrics.shoulder_height_diff,
                    'front_back_diff': metrics.shoulder_front_back_diff,
                    'narrowing': metrics.shoulder_narrowing,
                    'symmetry': metrics.shoulder_symmetry,
                    'status': self._get_shoulder_status(metrics)
                },
                'sitting_stability': {
                    'score': metrics.sitting_stability,
                    'status': self._get_metric_status(metrics.sitting_stability, 'sitting_stability')
                }
            },
            'suggestions': metrics.suggestions
        }
    
    def _get_metric_status(self, value: float, metric_type: str) -> str:
        """获取指标状态"""
        thresholds = self.thresholds[metric_type]
        
        if thresholds['excellent'][0] <= value <= thresholds['excellent'][1]:
            return 'excellent'
        elif thresholds['good'][0] <= value <= thresholds['good'][1]:
            return 'good'
        elif thresholds['fair'][0] <= value <= thresholds['fair'][1]:
            return 'fair'
        else:
            return 'poor'
    
    def _get_neck_status(self, metrics: PostureMetrics) -> str:
        """获取颈椎姿态状态"""
        forward_status = self._get_angle_status(metrics.neck_forward_angle, 'neck_posture', 'forward_angle')
        side_status = self._get_angle_status(metrics.neck_side_tilt, 'neck_posture', 'side_tilt')
        twist_status = self._get_angle_status(metrics.neck_twist_ratio, 'neck_posture', 'twist_ratio')
        
        # 取最差的状态
        statuses = [forward_status, side_status, twist_status]
        if 'poor' in statuses:
            return 'poor'
        elif 'fair' in statuses:
            return 'fair'
        elif 'good' in statuses:
            return 'good'
        else:
            return 'excellent'
    
    def _get_shoulder_status(self, metrics: PostureMetrics) -> str:
        """获取肩膀姿态状态"""
        height_status = self._get_angle_status(metrics.shoulder_height_diff, 'shoulder_posture', 'height_diff')
        front_back_status = self._get_angle_status(metrics.shoulder_front_back_diff, 'shoulder_posture', 'front_back_diff')
        symmetry_status = self._get_metric_status(metrics.shoulder_symmetry, 'shoulder_posture')
        
        # 取最差的状态
        statuses = [height_status, front_back_status, symmetry_status]
        if 'poor' in statuses:
            return 'poor'
        elif 'fair' in statuses:
            return 'fair'
        elif 'good' in statuses:
            return 'good'
        else:
            return 'excellent'
    
    def _get_angle_status(self, value: float, category: str, subcategory: str) -> str:
        """获取角度类型指标的状态"""
        thresholds = self.thresholds[category][subcategory]
        
        if thresholds['excellent'][0] <= value <= thresholds['excellent'][1]:
            return 'excellent'
        elif thresholds['good'][0] <= value <= thresholds['good'][1]:
            return 'good'
        elif thresholds['fair'][0] <= value <= thresholds['fair'][1]:
            return 'fair'
        else:
            return 'poor' 

    def _get_adaptive_grade(self, score: float) -> str:
        """
        使用自适应阈值确定等级
        
        Args:
            score: 综合评分
        
        Returns:
            等级字符串
        """
        # 获取自适应阈值
        adaptive_thresholds = self.user_profile.get_adaptive_thresholds()
        strictness_factor = adaptive_thresholds.get('strictness_factor', 1.0)
        
        # 调整阈值
        excellent_threshold = 90 * strictness_factor
        good_threshold = 75 * strictness_factor
        fair_threshold = 60 * strictness_factor
        
        if score >= excellent_threshold:
            return "excellent"
        elif score >= good_threshold:
            return "good"
        elif score >= fair_threshold:
            return "fair"
        else:
            return "poor"
    
    def get_user_statistics(self) -> Dict:
        """获取用户统计信息"""
        if not self.user_profile.posture_history:
            return {}
        
        scores = [record['overall_score'] for record in self.user_profile.posture_history]
        
        return {
            'total_sessions': len(self.user_profile.posture_history),
            'average_score': np.mean(scores),
            'best_score': max(scores),
            'worst_score': min(scores),
            'recent_trend': self._calculate_trend(scores[-10:]) if len(scores) >= 10 else 0,
            'improvement_rate': self._calculate_improvement_rate(scores)
        }
    
    def _calculate_trend(self, scores: List[float]) -> float:
        """计算评分趋势"""
        if len(scores) < 2:
            return 0
        
        x = np.arange(len(scores))
        slope = np.polyfit(x, scores, 1)[0]
        return slope
    
    def _calculate_improvement_rate(self, scores: List[float]) -> float:
        """计算改善率"""
        if len(scores) < 20:
            return 0
        
        early_avg = np.mean(scores[:10])
        recent_avg = np.mean(scores[-10:])
        
        if early_avg > 0:
            return (recent_avg - early_avg) / early_avg * 100
        else:
            return 0
    
    def enable_learning(self, enable: bool = True):
        """启用/禁用学习功能"""
        self.learning_enabled = enable
        print(f"学习功能已{'启用' if enable else '禁用'}")
    
    def reset_user_profile(self):
        """重置用户档案"""
        self.user_profile.posture_history = []
        self.user_profile.custom_thresholds = {}
        self.user_profile.preference_weights = {}
        self.user_profile.save_profile()
        print("用户档案已重置") 